//========================================================================
//Copyright 2007-2010 David Yu dyuproject@gmail.com
//------------------------------------------------------------------------
//Licensed under the Apache License, Version 2.0 (the "License");
//you may not use this file except in compliance with the License.
//You may obtain a copy of the License at 
//http://www.apache.org/licenses/LICENSE-2.0
//Unless required by applicable law or agreed to in writing, software
//distributed under the License is distributed on an "AS IS" BASIS,
//WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//See the License for the specific language governing permissions and
//limitations under the License.
//========================================================================

package com.dyuproject.kvstore.bdb;

import static com.dyuproject.kvstore.Serializers.STRING;

import java.util.Collections;
import java.util.Comparator;
import java.util.List;
import java.util.Map.Entry;

import com.dyuproject.kvstore.AbstractStore;
import com.dyuproject.kvstore.Mutable;
import com.dyuproject.kvstore.Serializer;
import com.dyuproject.kvstore.StoreException;
import com.google.common.base.Function;
import com.google.common.base.Preconditions;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.sleepycat.je.Cursor;
import com.sleepycat.je.Database;
import com.sleepycat.je.DatabaseEntry;
import com.sleepycat.je.DatabaseException;
import com.sleepycat.je.Environment;
import com.sleepycat.je.LockConflictException;
import com.sleepycat.je.LockMode;
import com.sleepycat.je.OperationStatus;
import com.sleepycat.je.Transaction;

public final class BDBStore<T> extends AbstractStore<T>
{
    
    public static final int DEFAULT_MAX_RETRIES = Integer.getInteger("bdbstore.max_retries", 3);

    public static final Comparator<byte[]> BDB_JE_COMPARATOR = new Comparator<byte[]>()
    {
        public int compare(byte[] o1, byte[] o2)
        {
            return compareBytes(o1, 0, o1.length, o2, 0, o2.length);
        }
    };

    static void close(Cursor cursor, boolean throwExceptionOnFail)
    {
        try
        {
            cursor.close();
        }
        catch (DatabaseException e)
        {
            if (throwExceptionOnFail)
                throw new StoreException(e.getMessage(),e);
        }
    }

    static void abort(Transaction tx, boolean throwExceptionOnFail)
    {
        try
        {
            tx.abort();
        }
        catch (DatabaseException e)
        {
            if (throwExceptionOnFail)
                throw new StoreException(e.getMessage(),e);
        }
    }

    public static int compareBytes(byte[] data1, int offset1, int size1, byte[] data2, 
            int offset2, int size2)
    {
        for (int i = 0; i < size1 && i < size2; i++)
        {

            int b1 = 0xFF & data1[offset1 + i];
            int b2 = 0xFF & data2[offset2 + i];
            if (b1 < b2)
                return -1;
            else if (b1 > b2)
                return 1;
        }

        if (size1 < size2)
            return -1;
        else if (size1 > size2)
            return 1;
        else
            return 0;
    }
    
    final Environment environment;
    final Database database;
    final Comparator<byte[]> comparator;
    final int maxRetries;
    
    public BDBStore(Serializer<T> serializer, Environment environment, Database database)
    {
        this(serializer, environment, database, DEFAULT_MAX_RETRIES);
    }

    public BDBStore(Serializer<T> serializer, Environment environment, Database database, 
            int maxRetries)
    {
        super(serializer);
        Preconditions.checkNotNull(environment);
        Preconditions.checkNotNull(database);
        this.environment = environment;
        this.database = database;
        this.maxRetries = maxRetries;
        try
        {
            if (database.getConfig().getSortedDuplicates())
                throw new IllegalStateException("sorted duplicates not allowed.");
            Comparator<byte[]> c = database.getConfig().getBtreeComparator();
            this.comparator = c == null?BDB_JE_COMPARATOR:c;
        }
        catch (DatabaseException e)
        {
            throw new RuntimeException(e);
        }
    }

    public void put(String key, T value)
    {
        Preconditions.checkNotNull(key);
        Preconditions.checkNotNull(value);
        boolean success = false;
        Transaction tx = null;
        try
        {
            tx = environment.beginTransaction(null, null);

            DatabaseEntry keyEntry = new DatabaseEntry(STRING.ser(key));
            DatabaseEntry valueEntry = new DatabaseEntry(serializer.ser(value));
            
            database.put(tx, keyEntry, valueEntry);
            
            tx.commit();
            success = true;
        }
        catch (DatabaseException e)
        {
            throw new StoreException(e.getMessage(),e);
        }
        finally
        {
            if(!success && tx!=null)
                abort(tx, false);
        }
    }

    public void delete(String key)
    {
        Preconditions.checkNotNull(key);
        boolean success = false;
        Transaction tx = null;
        try
        {
            tx = environment.beginTransaction(null, null);
            DatabaseEntry keyEntry = new DatabaseEntry(STRING.ser(key));
            database.delete(tx, keyEntry);
            
            tx.commit();
            success = true;
        }
        catch (DatabaseException e)
        {
            throw new StoreException(e.getMessage(),e);
        }
        finally
        {
            if (!success && tx!=null)
                abort(tx, false);
        }
    }

    public void delete(List<String> keys)
    {
        Preconditions.checkNotNull(keys);
        if (keys.size() == 0)
            return;
        boolean success = false;
        Transaction tx = null;
        try
        {
            tx = environment.beginTransaction(null, null);
            Cursor cursor = database.openCursor(tx, null);
            
            try
            {
                DatabaseEntry keyEntry = new DatabaseEntry();
                DatabaseEntry valueEntry = new DatabaseEntry();
                for (String k : keys)
                {
                    if (k != null)
                    {
                        keyEntry.setData(STRING.ser(k));
                        if (OperationStatus.SUCCESS == cursor.getSearchKey(keyEntry, valueEntry, 
                                LockMode.RMW))
                        {
                            cursor.delete();
                        }
                    }
                }
            }
            finally
            {
                cursor.close();
            }

            tx.commit();
            success = true;
        }
        catch (DatabaseException e)
        {
            throw new StoreException(e.getMessage(),e);
        }
        finally
        {
            if(!success && tx!=null)
                abort(tx, false);
        }
    }
    
    public T mutate(String key, Function<? super T, ? extends T> mutator)
    {
        return mutate(key, mutator, maxRetries);
    }

    public T mutate(String key, Function<? super T, ? extends T> mutator, int retries)
    {
        Preconditions.checkNotNull(key);
        Preconditions.checkNotNull(mutator);
        boolean success = false;
        T data = null;
        DatabaseException lastException = null;
        
        DatabaseEntry keyEntry = new DatabaseEntry(STRING.ser(key));
        DatabaseEntry valueEntry = new DatabaseEntry();
        
        for(int i=0; i<retries; i++)
        {
            Transaction tx = null;
            try
            {
                tx = environment.beginTransaction(null, null);
                Cursor cursor = database.openCursor(tx, null);
                
                try
                {
                    if (OperationStatus.SUCCESS == cursor.getSearchKey(keyEntry, valueEntry, 
                            LockMode.RMW))
                    {
                        data = mutator.apply(serializer.deser(valueEntry.getData()));
                        if (data == null)
                            cursor.delete();
                        else
                        {
                            valueEntry.setData(serializer.ser(data));
                            cursor.putCurrent(valueEntry);
                        }
                    }
                    else
                    {
                        data = mutator.apply(null);
                        if (data != null)
                        {
                            valueEntry.setData(serializer.ser(data));
                            cursor.put(keyEntry, valueEntry);
                        }
                    }
                }
                finally
                {
                    cursor.close();
                }
                
                tx.commit();
                success = true;
                return data;
            }
            catch(LockConflictException e)
            {
                lastException = e;
                continue;
            }
            catch (DatabaseException e)
            {
                lastException = e;
                continue;
                // break instead?
            }
            finally
            {
                if(!success && tx!=null)
                    abort(tx, false);
            }
        }

        throw new StoreException("Could not store data for key " + key,lastException);
    }
    
    public T mutateRaw(Mutable<String> mutableKey, Function<? super byte[], ? extends T> mutator)
    {
        return mutateRaw(mutableKey, mutator, maxRetries);
    }

    public T mutateRaw(Mutable<String> mutableKey, Function<? super byte[], ? extends T> mutator, 
            int retries)
    {
        Preconditions.checkNotNull(mutableKey);
        Preconditions.checkNotNull(mutator);
        String key = mutableKey.getValue();
        Preconditions.checkNotNull(key);
        
        boolean success = false;
        T data = null;
        DatabaseException lastException = null;
        
        DatabaseEntry keyEntry = new DatabaseEntry(STRING.ser(key));
        DatabaseEntry valueEntry = new DatabaseEntry();
        
        for(int i=0; i<retries; i++)
        {
            Transaction tx = null;
            try
            {
                tx = environment.beginTransaction(null, null);
                Cursor cursor = database.openCursor(tx, null);
                
                try
                {
                    if (OperationStatus.SUCCESS == cursor.getSearchKey(keyEntry, valueEntry, 
                            LockMode.RMW))
                    {
                        data = mutator.apply(valueEntry.getData());
                        String checkKey = mutableKey.getValue();
                        // abort
                        if(checkKey == null)
                            return data;
                        
                        if(checkKey != key)
                        {
                            key = checkKey;
                            keyEntry.setData(STRING.ser(checkKey));
                            continue;
                        }
                        
                        if (data == null)
                            cursor.delete();
                        else
                        {
                            valueEntry.setData(serializer.ser(data));
                            cursor.putCurrent(valueEntry);
                        }
                    }
                    else
                    {
                        data = mutator.apply(null);
                        String checkKey = mutableKey.getValue();
                        // abort
                        if(checkKey == null)
                            return data;
                        
                        if(checkKey != key)
                        {
                            key = checkKey;
                            keyEntry.setData(STRING.ser(checkKey));
                            continue;
                        }
                        
                        if (data != null)
                        {
                            valueEntry.setData(serializer.ser(data));
                            cursor.put(keyEntry, valueEntry);
                        }
                    }
                }
                finally
                {
                    cursor.close();
                }
                
                tx.commit();
                success = true;
                return data;
            }
            catch(LockConflictException e)
            {
                lastException = e;
                continue;
            }
            catch (DatabaseException e)
            {
                lastException = e;
                continue;
                // break instead?
            }
            finally
            {
                if(!success && tx!=null)
                    abort(tx, false);
            }
        }

        throw new StoreException("Could not store data for key " + key,lastException);
    }
    
    public <V> V get(String key, Serializer<V> serializer)
    {
        Preconditions.checkNotNull(key);
        V data = null;
        try
        {
            DatabaseEntry keyEntry = new DatabaseEntry(STRING.ser(key));
            DatabaseEntry valueEntry = new DatabaseEntry();

            if (OperationStatus.SUCCESS == database.get(null, keyEntry, valueEntry, 
                    LockMode.READ_UNCOMMITTED))
            {
                data = serializer.deser(valueEntry.getData());
            }
        }
        catch (DatabaseException e)
        {
            throw new StoreException(e.getMessage(),e);
        }
        return data;
    }
    
    public <V> List<V> get(List<String> keys, Serializer<V> serializer)
    {
        Preconditions.checkNotNull(keys);
        if (keys.size() == 0)
            return Collections.emptyList();

        List<V> result = Lists.newArrayListWithCapacity(keys.size());
        Cursor cursor = null;
        try
        {
            cursor = database.openCursor(null, null);
            
            DatabaseEntry keyEntry = new DatabaseEntry();
            DatabaseEntry valueEntry = new DatabaseEntry();
            for (String k : keys)
            {
                if (k != null)
                {
                    keyEntry.setData(STRING.ser(k));
                    if (OperationStatus.SUCCESS == cursor.getSearchKey(keyEntry, valueEntry, 
                            LockMode.READ_UNCOMMITTED))
                    {
                        result.add(serializer.deser(valueEntry.getData()));
                    }
                }
            }
        }
        catch (DatabaseException e1)
        {
            throw new StoreException(e1.getMessage(),e1);
        }
        finally
        {
            if (cursor != null)
                close(cursor, false);
        }

        return result;
    }

    public <V> List<Entry<String, V>> getEntries(List<String> keys, Serializer<V> serializer)
    {
        Preconditions.checkNotNull(keys);
        if (keys.size() == 0)
            return Collections.emptyList();

        List<Entry<String, V>> result = Lists.newArrayListWithCapacity(keys.size());
        Cursor cursor = null;
        try
        {
            cursor = database.openCursor(null, null);
            
            DatabaseEntry keyEntry = new DatabaseEntry();
            DatabaseEntry valueEntry = new DatabaseEntry();
            for (String k : keys)
            {
                if (k != null)
                {
                    keyEntry.setData(STRING.ser(k));
                    if (OperationStatus.SUCCESS == cursor.getSearchKey(keyEntry, valueEntry, 
                            LockMode.READ_UNCOMMITTED))
                    {
                        result.add(Maps.immutableEntry(k, serializer.deser(valueEntry.getData())));
                    }
                }
            }
        }
        catch (DatabaseException e1)
        {
            throw new StoreException(e1.getMessage(),e1);
        }
        finally
        {
            if (cursor != null)
                close(cursor, false);
        }

        return result;
    }

    public <V> List<Entry<String, V>> scanEntries(String start, String end, int max, 
            Serializer<V> serializer)
    {
        List<Entry<String, V>> results = Collections.emptyList();
        if (max < 1)
            return results;

        Cursor cursor = null;
        try
        {
            cursor = database.openCursor(null, null);
            byte[] endCompare = STRING.ser(end);
            
            DatabaseEntry keyEntry = new DatabaseEntry(STRING.ser(start));
            DatabaseEntry valueEntry = new DatabaseEntry();
            
            if (OperationStatus.SUCCESS == cursor.getSearchKeyRange(keyEntry, valueEntry, 
                    LockMode.READ_UNCOMMITTED)
                    && comparator.compare(keyEntry.getData(), endCompare) < 0)
            {
                results = Lists.newArrayList();
                results.add(Maps.immutableEntry(STRING.deser(keyEntry.getData()), 
                        serializer.deser(valueEntry.getData())));

                while (--max > 0 && OperationStatus.SUCCESS == cursor.getNext(keyEntry, 
                        valueEntry, LockMode.READ_UNCOMMITTED)
                        && comparator.compare(keyEntry.getData(), endCompare) < 0)
                {
                    results.add(Maps.immutableEntry(STRING.deser(keyEntry.getData()), 
                            serializer.deser(valueEntry.getData())));
                }
            }
        }
        catch (DatabaseException e)
        {
            throw new StoreException(e.getMessage(), e);
        }
        finally
        {
            if (cursor != null)
                close(cursor, false);
        }
        return results;
    }
    

    public <V> List<Entry<String, V>> scanEntriesReverse(String start, String end, int max, 
            Serializer<V> serializer)
    {
        List<Entry<String, V>> results = Collections.emptyList();
        if (max < 1)
            return results;

        Cursor cursor = null;
        try
        {
            cursor = database.openCursor(null, null);
            byte[] startCompare = STRING.ser(start);
            
            DatabaseEntry keyEntry = new DatabaseEntry(STRING.ser(end));
            DatabaseEntry valueEntry = new DatabaseEntry();
            
            if (OperationStatus.SUCCESS == cursor.getSearchKeyRange(keyEntry, valueEntry, 
                    LockMode.READ_UNCOMMITTED))
            {
                results = Lists.newArrayList();
                while (max-- > 0 && OperationStatus.SUCCESS == cursor.getPrev(keyEntry, 
                        valueEntry, LockMode.READ_UNCOMMITTED)
                        && comparator.compare(keyEntry.getData(), startCompare) >= 0)
                {
                    results.add(Maps.immutableEntry(STRING.deser(keyEntry.getData()), 
                            serializer.deser(valueEntry.getData())));
                }
            }
            else if (OperationStatus.SUCCESS == cursor.getLast(keyEntry, valueEntry, 
                    LockMode.READ_UNCOMMITTED)
                    && comparator.compare(keyEntry.getData(), startCompare) >= 0)
            {
                results = Lists.newArrayList();
                results.add(Maps.immutableEntry(STRING.deser(keyEntry.getData()), 
                        serializer.deser(valueEntry.getData())));
                while (--max > 0 && OperationStatus.SUCCESS == cursor.getPrev(keyEntry, 
                        valueEntry, LockMode.READ_UNCOMMITTED)
                        && comparator.compare(keyEntry.getData(), startCompare) >= 0)
                {
                    results.add(Maps.immutableEntry(STRING.deser(keyEntry.getData()), 
                            serializer.deser(valueEntry.getData())));
                }
            }
        }
        catch (DatabaseException e)
        {
            throw new StoreException(e.getMessage(), e);
        }
        finally
        {
            if (cursor != null)
                close(cursor, false);
        }
        return results;
    }
    
    public <V> List<V> scan(String start, String end, int max, Serializer<V> serializer)
    {
        List<V> results = Collections.emptyList();
        if (max < 1)
            return results;

        Cursor cursor = null;
        try
        {
            cursor = database.openCursor(null, null);
            byte[] endCompare = STRING.ser(end);
            
            DatabaseEntry keyEntry = new DatabaseEntry(STRING.ser(start));
            DatabaseEntry valueEntry = new DatabaseEntry();
            
            if (OperationStatus.SUCCESS == cursor.getSearchKeyRange(keyEntry, valueEntry, 
                    LockMode.READ_UNCOMMITTED)
                    && comparator.compare(keyEntry.getData(), endCompare) < 0)
            {
                results = Lists.newArrayList();
                results.add(serializer.deser(valueEntry.getData()));

                while (--max > 0 && OperationStatus.SUCCESS == cursor.getNext(keyEntry, 
                        valueEntry, LockMode.READ_UNCOMMITTED)
                        && comparator.compare(keyEntry.getData(), endCompare) < 0)
                {
                    results.add(serializer.deser(valueEntry.getData()));
                }
            }
        }
        catch (DatabaseException e)
        {
            throw new StoreException(e.getMessage(), e);
        }
        finally
        {
            if (cursor != null)
                close(cursor, false);
        }
        return results;
    }
    
    public <V> List<V> scanReverse(String start, String end, int max, Serializer<V> serializer)
    {
        List<V> results = Collections.emptyList();
        if (max < 1)
            return results;

        Cursor cursor = null;
        try
        {
            cursor = database.openCursor(null, null);
            byte[] startCompare = STRING.ser(start);
            
            DatabaseEntry keyEntry = new DatabaseEntry(STRING.ser(end));
            DatabaseEntry valueEntry = new DatabaseEntry();
            
            if (OperationStatus.SUCCESS == cursor.getSearchKeyRange(keyEntry, valueEntry, 
                    LockMode.READ_UNCOMMITTED))
            {
                results = Lists.newArrayList();
                while (max-- > 0 && OperationStatus.SUCCESS == cursor.getPrev(keyEntry, 
                        valueEntry, LockMode.READ_UNCOMMITTED)
                        && comparator.compare(keyEntry.getData(), startCompare) >= 0)
                {
                    results.add(serializer.deser(valueEntry.getData()));
                }
            }
            else if (OperationStatus.SUCCESS == cursor.getLast(keyEntry, valueEntry, 
                    LockMode.READ_UNCOMMITTED)
                    && comparator.compare(keyEntry.getData(), startCompare) >= 0)
            {
                results = Lists.newArrayList();
                results.add(serializer.deser(valueEntry.getData()));
                while (--max > 0 && OperationStatus.SUCCESS == cursor.getPrev(keyEntry, 
                        valueEntry, LockMode.READ_UNCOMMITTED)
                        && comparator.compare(keyEntry.getData(), startCompare) >= 0)
                {
                    results.add(serializer.deser(valueEntry.getData()));
                }
            }
        }
        catch (DatabaseException e)
        {
            throw new StoreException(e.getMessage(), e);
        }
        finally
        {
            if (cursor != null)
                close(cursor, false);
        }
        return results;
    }

    public List<String> keyScan(String start, String end, int max)
    {
        List<String> results = Collections.emptyList();
        if (max < 1)
            return results;

        Cursor cursor = null;
        try
        {
            cursor = database.openCursor(null, null);
            byte[] endCompare = STRING.ser(end);
            
            DatabaseEntry keyEntry = new DatabaseEntry(STRING.ser(start));
            DatabaseEntry valueEntry = new DatabaseEntry();
            
            if (OperationStatus.SUCCESS == cursor.getSearchKeyRange(keyEntry, valueEntry, 
                    LockMode.READ_UNCOMMITTED)
                    && comparator.compare(keyEntry.getData(), endCompare) < 0)
            {
                results = Lists.newArrayList();
                results.add(STRING.deser(keyEntry.getData()));

                while (--max > 0 && OperationStatus.SUCCESS == cursor.getNext(keyEntry, 
                        valueEntry, LockMode.READ_UNCOMMITTED)
                        && comparator.compare(keyEntry.getData(), endCompare) < 0)
                {
                    results.add(STRING.deser(keyEntry.getData()));
                }
            }
        }
        catch (DatabaseException e)
        {
            throw new StoreException(e.getMessage(), e);
        }
        finally
        {
            if (cursor != null)
                close(cursor, false);
        }
        return results;
    }

    public List<String> keyScanReverse(String start, String end, int max)
    {
        List<String> results = Collections.emptyList();
        if (max < 1)
            return results;

        Cursor cursor = null;
        try
        {
            cursor = database.openCursor(null, null);
            byte[] startCompare = STRING.ser(start);
            
            DatabaseEntry keyEntry = new DatabaseEntry(STRING.ser(end));
            DatabaseEntry valueEntry = new DatabaseEntry();
            
            if (OperationStatus.SUCCESS == cursor.getSearchKeyRange(keyEntry, valueEntry, 
                    LockMode.READ_UNCOMMITTED))
            {
                results = Lists.newArrayList();
                while (max-- > 0 && OperationStatus.SUCCESS == cursor.getPrev(keyEntry, valueEntry, LockMode.READ_UNCOMMITTED)
                        && comparator.compare(keyEntry.getData(), startCompare) >= 0)
                {
                    results.add(STRING.deser(keyEntry.getData()));
                }
            }
            else if (OperationStatus.SUCCESS == cursor.getLast(keyEntry, valueEntry, 
                    LockMode.READ_UNCOMMITTED)
                    && comparator.compare(keyEntry.getData(), startCompare) >= 0)
            {
                results = Lists.newArrayList();
                results.add(STRING.deser(keyEntry.getData()));
                while (--max > 0 && OperationStatus.SUCCESS == cursor.getPrev(keyEntry, 
                        valueEntry, LockMode.READ_UNCOMMITTED)
                        && comparator.compare(keyEntry.getData(), startCompare) >= 0)
                {
                    results.add(STRING.deser(keyEntry.getData()));
                }
            }
        }
        catch (DatabaseException e)
        {
            throw new StoreException(e.getMessage(), e);
        }
        finally
        {
            if (cursor != null)
                close(cursor, false);
        }
        return results;
    }

}
